from chromadb import PersistentClient
from pathlib import Path
import json
import requests
from config import OLLAMA_BASE_URL, EMBEDDING_MODEL, CHROMA_CONFIG
from chromadb.config import Settings

class ChromaService:
    # settings=Settings(anonymized_telemetry=False)
    def __init__(self):
        self.client = PersistentClient(path=str(Path(CHROMA_CONFIG['path']).resolve()), settings=Settings(anonymized_telemetry=False))
    
    def list_collections(self, page, page_size, search):
        collection_names = self.client.list_collections()
        collections = []
        
        # 获取所有collections的信息
        for name in collection_names:
            # 如果有搜索条件，先过滤
            if search and search not in name.lower():
                continue
                
            col = self.client.get_collection(name)
            result = col.get()
            dimension = None
            if result.get("embeddings") and len(result["embeddings"]) > 0:
                dimension = len(result["embeddings"][0])
            
            collections.append({
                'name': col.name,
                'dimension': dimension,
                'count': col.count()
            })
        
        # 计算总数和分页
        total = len(collections)
        start = (page - 1) * page_size
        end = start + page_size
        
        return {
            'collections': collections[start:end],
            'total': total,
            'page': page,
            'page_size': page_size
        }
    
    def create_collection(self, name):
        # if self.client.has_collection(name):
        #     raise Exception("Collection already exists")
        
        collection = self.client.create_collection(name)
        return {'message': 'Collection created', 'collection': {'name': collection.name}}
    
    def update_collection(self, collection_name, new_name):
        # 获取旧集合
        old_collection = self.client.get_collection(name=collection_name)
        # 新建新集合
        new_collection = self.client.create_collection(new_name)
        # 复制旧集合中的所有记录
        result = old_collection.get()
        ids = result.get("ids", [])
        embeddings = result.get("embeddings", [])
        metadatas = result.get("metadatas", [])
        documents = result.get("documents", [])
        if ids:
            new_collection.add(ids=ids, embeddings=embeddings, metadatas=metadatas, documents=documents)
        # 删除旧集合
        self.client.delete_collection(collection_name)
        return {"message": "Collection renamed successfully", "collection": {"name": new_name}}
    
    def delete_collection(self, collection_name):
        self.client.delete_collection(collection_name)
        return {"message": "Collection deleted successfully"}
    
    def list_records(self, collection_name, page, page_size, search_type, record_id, query_text, metadata_filter, document_filter, n_results_requested):
        collection = self.client.get_collection(name=collection_name)
        
        # 处理元数据过滤条件
        where_clause = None
        if metadata_filter:
            try:
                # 如果是字符串，尝试解析为字典
                if isinstance(metadata_filter, str):
                    filter_dict = json.loads(metadata_filter)
                # 如果已经是字典，直接使用
                elif isinstance(metadata_filter, dict):
                    filter_dict = metadata_filter
                else:
                    raise ValueError("metadata_filter must be either a string or a dictionary")

                if isinstance(filter_dict, dict):
                    # 如果是多个条件，转换为 $and 格式
                    if len(filter_dict) > 1:
                        where_clause = {
                            "$and": [
                                {key: {"$eq": value}} if isinstance(value, (str, int, float, bool)) else {key: value}
                                for key, value in filter_dict.items()
                            ]
                        }
                    # 单个条件，直接转换
                    else:
                        key, value = next(iter(filter_dict.items()))
                        where_clause = {
                            key: {"$eq": value} if isinstance(value, (str, int, float, bool)) else value
                        }
            except json.JSONDecodeError:
                return {"error": "Invalid metadata filter format"}, 400

        # ID 精确查询
        if search_type == 'id' and record_id:
            result = collection.get(
                ids=[record_id],
                include=["metadatas", "documents"]
            )
            total = len(result["ids"])
        
        # 语义查询
        elif search_type == 'hybrid' and query_text:
            try:
                # 使用 Ollama 生成查询文本的 embedding
                query_embedding = self.get_embedding(query_text)
                
                # 获取查询参数
                query_params = {
                    "query_embeddings": [query_embedding],
                    "n_results": n_results_requested,
                    "include": ["metadatas", "documents", "distances"]
                }
                if where_clause:
                    query_params["where"] = where_clause
                if document_filter:
                    query_params["where_document"] = {"$contains": document_filter}
                
                result = collection.query(**query_params)
                
                # 处理嵌套数组结果
                result = {
                    "ids": result["ids"][0] if isinstance(result["ids"][0], list) else result["ids"],
                    "documents": result["documents"][0] if isinstance(result["documents"][0], list) else result["documents"],
                    "metadatas": result["metadatas"][0] if result.get("metadatas") and isinstance(result["metadatas"][0], list) else result.get("metadatas", []),
                    "distances": result["distances"][0] if result.get("distances") and isinstance(result["distances"][0], list) else result.get("distances", [])
                }
                # 手动处理分页
                start_idx = (page - 1) * page_size
                end_idx = start_idx + page_size
                result = {
                    "ids": result["ids"][start_idx:end_idx],
                    "documents": result["documents"][start_idx:end_idx],
                    "metadatas": result["metadatas"][start_idx:end_idx] if result.get("metadatas") else [],
                    "distances": result["distances"][start_idx:end_idx] if result.get("distances") else []
                }
                # 计算符合过滤条件的总记录数，并取 n_results_requested 与过滤总数的最小值
                total_count = self.get_filtered_count(
                    collection,
                    where=where_clause,
                    where_document={"$contains": document_filter} if document_filter else None
                )
                total = min(total_count, n_results_requested)
            except Exception as e:
                return {"error": str(e)}, 500
        
        # 普通查询（包括元数据和文档内容过滤）
        else:
            # 使用 get 方法的内置分页
            get_params = {
                "limit": page_size,
                "offset": (page - 1) * page_size,
                "include": ["metadatas", "documents"]
            }
            if where_clause:
                get_params["where"] = where_clause
            if document_filter:
                get_params["where_document"] = {"$contains": document_filter}
                
            result = collection.get(**get_params)
            if where_clause or document_filter:
                # 获取符合条件的总记录数
                total = self.get_filtered_count(
                    collection,
                    where=where_clause,
                    where_document={"$contains": document_filter} if document_filter else None
                )
            else:
                total = collection.count()  # 获取总记录数用于分页

        # 构建记录列表
        records = [
            {
                "id": doc_id,
                "document": doc,
                "metadata": meta if meta else {},
                "distance": dist if 'distances' in result else None
            }
            for doc_id, doc, meta, dist in zip(
                result["ids"],
                result["documents"],
                result.get("metadatas", [{}] * len(result["ids"])),
                result.get("distances", [None] * len(result["ids"]))
            )
        ]

        return {
            "records": records,
            "total": total,
            "page": page,
            "page_size": page_size
        }
    
    def create_record(self, collection_name, record_id, document, metadata):
        # 使用 Ollama 生成 embedding
        embedding = self.get_embedding(document)
        
        collection = self.client.get_collection(name=collection_name)
        
        # 基础参数
        add_params = {
            "ids": [record_id],
            "embeddings": [embedding],
            "documents": [document]
        }
        
        # 只有在 metadata 存在且不为空时才添加
        if metadata:
            add_params["metadatas"] = [metadata]
            
        collection.add(**add_params)
        return {
            "message": "Record added successfully", 
            "record_id": record_id
        }

    def update_record(self, collection_name, record_id, document, metadata):
        # 使用 Ollama 生成 embedding
        embedding = self.get_embedding(document)
        
        collection = self.client.get_collection(name=collection_name)
        # 先删除旧记录（如果存在）
        try:
            collection.delete(ids=[record_id])
        except Exception:
            pass
            
        # 基础参数
        add_params = {
            "ids": [record_id],
            "embeddings": [embedding],
            "documents": [document]
        }
        
        # 只有在 metadata 存在且不为空时才添加
        if metadata:
            add_params["metadatas"] = [metadata]
            
        collection.add(**add_params)
        return {
            "message": "Record updated successfully", 
            "record_id": record_id
        }

    def delete_record(self, collection_name, record_id):
        collection = self.client.get_collection(name=collection_name)
        collection.delete(ids=[record_id])
        return {
            "message": "Record deleted successfully", 
            "record_id": record_id
        }

    def batch_delete_records(self, collection_name, ids):
        collection = self.client.get_collection(name=collection_name)
        collection.delete(ids=ids)
        return {
            "message": "Records deleted successfully", 
            "ids": ids
        }

    def get_embedding(self, text):
        """使用 Ollama API 获取文本的 embedding"""
        try:
            response = requests.post(
                f"{OLLAMA_BASE_URL}/api/embeddings",
                json={
                    "model": EMBEDDING_MODEL,
                    "prompt": text
                }
            )
            response.raise_for_status()
            return response.json()["embedding"]
        except Exception as e:
            print(f"Error getting embedding: {str(e)}")
            raise

    def get_filtered_count(self, collection, where=None, where_document=None):
        """获取符合过滤条件的记录总数"""
        try:
            # 只获取 ids 以提高性能
            result = collection.get(
                where=where,
                where_document=where_document,
                include=[]  # 只包含 ids，不包含其他数据以提高性能
            )
            return len(result["ids"])
        except Exception as e:
            print(f"Error getting filtered count: {str(e)}")
            return 0
        
    def get_count(self, collection_name: str):
        """获取记录总数"""
        try:
            collection = self.client.get_collection(name=collection_name)
            # 只获取 ids 以提高性能
            result = collection.get(
                include=[]  # 只包含 ids，不包含其他数据以提高性能
            )
            return len(result["ids"])
        except Exception as e:
            print(f"Error getting filtered count: {str(e)}")
            return 0
    def delete_records_by_document_id(self, collection_name, document_id):
        """
        根据 metadata 中 document_id 字段删除所有记录
        参数:
            collection_name: 集合名称
            document_id: 要删除的 document_id
        返回:
            包含删除记录信息的字典
        """
        collection = self.client.get_collection(name=collection_name)
        # 构造过滤条件：metadata 中 document_id 等于给定 document_id
        where_clause = {
            "document_id": {"$eq": document_id}
        }
        print(where_clause)
        # 查询符合条件的记录（只获取 ids 提高效率）
        result = collection.get(where=where_clause, include=[])
        record_ids = result.get("ids", [])
        if record_ids:
            collection.delete(ids=record_ids)
        return {
            "message": "Records deleted successfully",
            "deleted_ids": record_ids,
            "document_id": document_id
        }